from backbones.rnn import RNN
import torch
from dataset.dataloader import generate_loader_vocab
import torch.nn.functional as F

if __name__ == '__main__':
    _, vocab = generate_loader_vocab(batch_size=10)
    vocab_size = len(vocab)
    hidden_size = 1000
    device = torch.device("cuda")
    model = RNN(vocab_size, hidden_size, num_layers=2, device=device)
    model.load_state_dict(torch.load("./save/best.pt", weights_only=True))

    input_str = torch.tensor(vocab.to_idx("a real traveller amid".split()), device=device)
    print(input_str)
    input_str = input_str.unsqueeze(-1)
    input_str = F.one_hot(input_str, vocab_size).float()

    h0 = torch.zeros((2, 1, hidden_size), device=device)
    model.eval()
    y_hat, h = model(input_str, h0)

    indices = torch.argmax(y_hat, dim=-1).squeeze(1)

    print([vocab[index.item()] for index in indices ])